"""AMD GPU 功耗采样后端（通过 ROCm SMI）"""

from __future__ import annotations
from typing import Optional
import subprocess
import re

from .base import PowerBackend


class ROCmBackend(PowerBackend):
    """AMD GPU 功耗采样（通过 ROCm SMI）"""
    
    def __init__(self):
        self._device_id = 0
    
    @staticmethod
    def is_available() -> bool:
        try:
            result = subprocess.run(['rocm-smi', '--showpower'], 
                                    capture_output=True, timeout=2)
            return result.returncode == 0
        except:
            return False
    
    def init(self, device_id: int = 0) -> bool:
        self._device_id = device_id
        return self.is_available()
    
    def sample(self) -> Optional[float]:
        try:
            result = subprocess.run(
                ['rocm-smi', f'-d{self._device_id}', '--showpower'],
                capture_output=True, text=True, timeout=1
            )
            # 解析输出，提取功率值（单位：W）
            match = re.search(r'Average Graphics Package Power:\s+([\d.]+)\s*W', result.stdout)
            if match:
                power_W = float(match.group(1))
                return power_W * 1000.0  # 转换为 mW
            return None
        except:
            return None
    
    def cleanup(self):
        pass
    
    def get_name(self) -> str:
        return "rocm"

